from mind3d.dataset.kitti_eagermot import MOTDatasetKITTI
from mind3d.dataset.nuscenes_eagermot import MOTDatasetNuScenes
from mind3d.models.eagermot import Eagermot


def build_dataset(source, **kwargs):
    """build_dataset

    Args:
        source (str): whether kitti or nuscenes

    Returns:
        dataset: returned dataset
    """
    if source=="kitti":
        return MOTDatasetKITTI(det_source=kwargs['det_source'], work_dir=kwargs['work_dir'], 
                                seg_source=kwargs['seg_source'], data_dir=kwargs['data_dir'], param=kwargs['param'],
                                args=kwargs['args'])
    elif source=="nuscenes":
        return MOTDatasetNuScenes(det_source=kwargs['det_source'], work_dir=kwargs['work_dir'], 
                                seg_source=kwargs['seg_source'], data_dir=kwargs['data_dir'], param=kwargs['param'],
                                args=kwargs['args'], version=kwargs['version'])
    else:
        return 0
    

def build_model(source, **kwargs):
    """
    Args:
        source (str): model name eg. eagermot

    Returns:
        model: returned model
    """
    if source=="eagermot":
        return Eagermot(kwargs['sequence'], kwargs['params'], kwargs['test_mode'])
    else:
        return 0